/*
Copyright (c) 2025 Huawei Technologies Co., Ltd.
This file is a part of the CANN Open Software.
Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
Please refer to the License for details. You may not use this file except in compliance with the License.
THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
See LICENSE in the root of the software repository for the full text of the License. 
*/

#pragma once
#include "./custom_type.h"
#include "./data_transfer.h"
#include "./kernel_const.h"
#include "./kernel_utils.h"
#include "./swizzle.h"
#include "kernel_operator.h"


/**
 * @brief: 批量线性乘加，完成分核和UB缓存分配，实现 C = alpha * AB + beta * C
 * @param [in] L1_M0: 搬入L1计算的A矩阵M维度大小
 * @param [in] L1_N0: 搬入L1计算的B矩阵N维度大小
 * @param [in] L1_K0: 搬入L1计算的AB矩阵K维度大小
 * @param [in] WORKSPACE_NUM: 显存开辟的workspace数
 * @param [in] zeroPaddingM: A、C矩阵零填充后的M维度
 * @param [in] zeroPaddingN: B、C矩阵零填充后的N维度
 * @param [in] batchCount: 批量矩阵乘的batch数
 * @param [in] d_validM: 每批矩阵乘的A、C矩阵M维度有效长度数组
 * @param [in] d_validN: 每批矩阵乘的B、C矩阵N维度有效长度数组
 * @param [in] alpha: alpha*AB+beta*C
 * @param [in] beta:  alpha*AB+beta*C
 * @param [out] d_CPointer: 每批矩阵乘的零填充B矩阵首地址数组
 * @param [out] d_AicAivWorkspacePointer: Aic Aiv 同步的GM空间首地址数组
 * @param [in] isAlpha1Beta0: 是否有 alpha==1.0 && beta==0.0
 */
 template<
    uint32_t L1_M0, 
    uint32_t L1_N0, 
    uint32_t L1_K0,
    uint32_t WORKSPACE_NUM
>
[aicore] inline __attribute__((always_inline)) void BatchMatmulEpilogue(
    uint32_t zeroPaddingM, 
    uint32_t zeroPaddingN, 
    uint32_t batchCount, 
    __gm__ uint32_t*  d_validM,
    __gm__ uint32_t*  d_validN,
    half alpha,
    half beta,
    __gm__ half**  d_CPointer, 
    __gm__ half** d_AicAivWorkspacePointer, 
    uint8_t isAlpha1Beta0 
){

    if(isAlpha1Beta0){

        // no epilogue

    }else{ // epilogue

        AscendC::TBuf<AscendC::TPosition::VECIN> ub_buf;
        AscendC::TPipe ub_pipe;
        ub_pipe.InitBuffer(ub_buf, UB_BYTES); 
        AscendC::LocalTensor<uint8_t> ub_tensor = ub_buf.Get<uint8_t>(); 
        ub_pipe.Destroy();

        uint32_t groupNum = AscendC::GetBlockNum(); 
        uint32_t groupIdx = AscendC::GetBlockIdx()/AscendC::GetSubBlockNum(); 
        static constexpr uint32_t aivInnerNum = 2; // AscendC::GetSubBlockNum()
        uint32_t aivInnerIdx = AscendC::GetSubBlockIdx(); 
        uint32_t aivTotalNum = AscendC::GetBlockNum() * AscendC::GetSubBlockNum(); 
        uint32_t aivTotalIdx = AscendC::GetBlockIdx(); 
         

        static constexpr uint32_t l0c_blockSize = L1_M0 * L1_N0; 

        static constexpr uint32_t ub_pingpongNum = UB_BYTES / 3 / (l0c_blockSize * sizeof(half) / aivInnerNum);
        AscendC::LocalTensor<half> ub_cBuf[ub_pingpongNum]; 
        AscendC::LocalTensor<half> ub_abBuf[ub_pingpongNum];
        AscendC::LocalTensor<half> ub_resultBuf[ub_pingpongNum];

        uint32_t ub_pingpongBytes = l0c_blockSize * sizeof(half) / 2;
        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            ub_cBuf[i] = ub_tensor[                             i * ub_pingpongBytes].template ReinterpretCast<half>();
        }
        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            ub_abBuf[i] = ub_tensor[     (ub_pingpongNum +     i) * ub_pingpongBytes].template ReinterpretCast<half>();
        }
        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            ub_resultBuf[i] = ub_tensor[ (ub_pingpongNum * 2 + i) * ub_pingpongBytes].template ReinterpretCast<half>();
        }

        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>((event_t)(i));
        }
        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>((event_t)(ub_pingpongNum + i));
        }

        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>((event_t)(2 * ub_pingpongNum + i));
        }

        AscendC::GlobalTensor<half> gm_workspace;   
        gm_workspace.SetGlobalBuffer( (__gm__ half*)d_AicAivWorkspacePointer[groupIdx] );
        AscendC::GlobalTensor<half> gm_workspacePingpongBuf[WORKSPACE_NUM];
        for(uint32_t i = 0; i < WORKSPACE_NUM; i++){
            gm_workspacePingpongBuf[i] = gm_workspace[ i * l0c_blockSize ];
        }

        for(uint32_t i = 0; i < WORKSPACE_NUM; i++){
            AscendC::CrossCoreSetFlag<0x2, PIPE_MTE2>(WORKSPACE_NUM + i);
        }

        AscendC::GlobalTensor<half> gm_tensorC;

        

        uint32_t M = 0; 
        uint32_t N = 0; 
        uint32_t strideC = 0; 
        uint32_t mLoops = 0;
        uint32_t nLoops = 0;

        uint32_t curAicoreBlock = 0;
        uint32_t curResultBlocksSum = 0;
        uint32_t batchNum = -1;
        int64_t loopSum = 0;

        for(uint32_t loopIdx = 0; loopIdx < loopSum+1; loopIdx++){

            if( loopIdx == loopSum ){
                batchNum++; 
                if( batchNum < batchCount ){
                    curResultBlocksSum = loopSum;
                    gm_tensorC.SetGlobalBuffer( (__gm__ half*)d_CPointer[batchNum] );
                    loopSum += CeilDiv<int64_t>(d_validM[batchNum], L1_M0) * CeilDiv<int64_t>(d_validN[batchNum], L1_N0);
                    M = d_validM[batchNum]; 
                    N = d_validN[batchNum]; 
                    strideC = zeroPaddingN; 
                    mLoops = CeilDiv<uint32_t>(M, L1_M0);
                    nLoops = CeilDiv<uint32_t>(N, L1_N0);
                }else{
                    continue;
                } 

            }

            if(loopIdx % groupNum != groupIdx){
                continue;
            }


            uint32_t blockIdx = loopIdx - curResultBlocksSum;
            uint32_t blockMIdx = blockIdx / nLoops;
            uint32_t blockNIdx = blockIdx % nLoops;
            BlockSwizzle(
                blockMIdx, 
                blockNIdx, 
                blockIdx, 
                mLoops, 
                nLoops
            ); 
            uint32_t gm_mActual = (blockMIdx == mLoops-1? M-blockMIdx*L1_M0 : L1_M0 );
            uint32_t gm_nActual = (blockNIdx == nLoops-1? N-blockNIdx*L1_N0 : L1_N0 );

            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>((event_t)(curAicoreBlock % ub_pingpongNum));

            for(uint32_t partIdx = 0; partIdx < aivInnerNum; partIdx++){
                if(aivInnerIdx != partIdx){
                    continue; 
                }
                uint32_t aivCalTaskLen = L1_M0 / aivInnerNum; 
                uint32_t preCalLen = aivInnerIdx * aivCalTaskLen; 
                if(gm_mActual > preCalLen){
                    uint32_t curCalLen = (gm_mActual - preCalLen < aivCalTaskLen ? gm_mActual - preCalLen : aivCalTaskLen); 
                    Gm2Ub<half>(
                        ub_cBuf[curAicoreBlock % ub_pingpongNum], 
                        gm_tensorC[(blockMIdx * L1_M0 + preCalLen) * strideC + blockNIdx * L1_N0], 
                        curCalLen, 
                        gm_nActual, 
                        0, 
                        strideC, 
                        0, 
                        0
                    );
                }
            }

            AscendC::SetFlag<AscendC::HardEvent::MTE2_V>((event_t)(curAicoreBlock % ub_pingpongNum));
            
            AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>((event_t)(curAicoreBlock % ub_pingpongNum));

            AscendC::Muls<half, false>(
                ub_cBuf[curAicoreBlock % ub_pingpongNum], 
                ub_cBuf[curAicoreBlock % ub_pingpongNum], 
                beta, 
                AscendC::MASK_PLACEHOLDER, 
                l0c_blockSize / 2 / (8 * DATABLOCK_BYTES / sizeof(half)), 
                AscendC::UnaryRepeatParams(1, 1, 8, 8)
            );

            AscendC::PipeBarrier<PIPE_V>();

            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>((event_t)(ub_pingpongNum + curAicoreBlock % ub_pingpongNum));

            AscendC::CrossCoreWaitFlag(curAicoreBlock % WORKSPACE_NUM);

            for(uint32_t partIdx = 0; partIdx < aivInnerNum; partIdx++){
                if(aivInnerIdx != partIdx){
                    continue; 
                }
                uint32_t aivCalTaskLen = L1_M0 / aivInnerNum; 
                uint32_t preCalLen = aivInnerIdx * aivCalTaskLen; 
                if(gm_mActual > preCalLen){
                    uint32_t curCalLen = (gm_mActual - preCalLen < aivCalTaskLen ? gm_mActual - preCalLen : aivCalTaskLen ); 
                    Gm2Ub<half>(
                        ub_abBuf[curAicoreBlock % ub_pingpongNum],
                        gm_workspacePingpongBuf[curAicoreBlock % WORKSPACE_NUM][preCalLen * L1_N0], 
                        curCalLen, 
                        gm_nActual, 
                        0, 
                        L1_N0, 
                        0, 
                        0
                    );
                }
            }


            AscendC::CrossCoreSetFlag<0x2, PIPE_MTE2>(WORKSPACE_NUM + curAicoreBlock % WORKSPACE_NUM);

            AscendC::SetFlag<AscendC::HardEvent::MTE2_V>((event_t)(ub_pingpongNum + curAicoreBlock % ub_pingpongNum));

            AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>((event_t)(ub_pingpongNum + curAicoreBlock % ub_pingpongNum));

            AscendC::Muls<half, false>(
                ub_abBuf[curAicoreBlock % ub_pingpongNum], 
                ub_abBuf[curAicoreBlock % ub_pingpongNum], 
                alpha, 
                AscendC::MASK_PLACEHOLDER, 
                l0c_blockSize / 2 / (8 * DATABLOCK_BYTES / sizeof(half)),
                AscendC::UnaryRepeatParams(1, 1, 8, 8)
            );

            AscendC::PipeBarrier<PIPE_V>();

            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>((event_t)(2 * ub_pingpongNum + curAicoreBlock % ub_pingpongNum));

            AscendC::Add<half, false>(
                ub_resultBuf[curAicoreBlock % ub_pingpongNum], 
                ub_abBuf[curAicoreBlock % ub_pingpongNum], 
                ub_cBuf[curAicoreBlock % ub_pingpongNum], 
                AscendC::MASK_PLACEHOLDER, 
                l0c_blockSize / 2 / (8 * DATABLOCK_BYTES / sizeof(half)),
                AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8)
            );
            AscendC::PipeBarrier<PIPE_V>();

            AscendC::SetFlag<AscendC::HardEvent::V_MTE3>((event_t)(2 * ub_pingpongNum + curAicoreBlock % ub_pingpongNum));

            
            AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>((event_t)(2 * ub_pingpongNum + curAicoreBlock % ub_pingpongNum));
            for(uint32_t partIdx = 0; partIdx < aivInnerNum; partIdx++){
                if(aivInnerIdx != partIdx){
                    continue; 
                }
                uint32_t aivCalTaskLen = L1_M0 / aivInnerNum; 
                uint32_t preCalLen = aivInnerIdx * aivCalTaskLen; 
                if(gm_mActual > preCalLen){
                    uint32_t curCalLen = (gm_mActual - preCalLen < aivCalTaskLen ? gm_mActual - preCalLen : aivCalTaskLen); 
                    Ub2Gm<half>(
                        gm_tensorC[(blockMIdx * L1_M0 + preCalLen) * strideC + blockNIdx * L1_N0], 
                        ub_resultBuf[curAicoreBlock % ub_pingpongNum], 
                        curCalLen, 
                        gm_nActual, 
                        strideC, 
                        0
                    ); 
                }
            }

            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>((event_t)(curAicoreBlock % ub_pingpongNum));
            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>((event_t)(ub_pingpongNum + curAicoreBlock % ub_pingpongNum));
            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>((event_t)(2 * ub_pingpongNum + curAicoreBlock % ub_pingpongNum));

            curAicoreBlock++;
        }

        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>((event_t)(i));
        }
        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>((event_t)(ub_pingpongNum + i));
        }
        #pragma unroll
        for(uint32_t i = 0; i < ub_pingpongNum; i++){
            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>((event_t)(2 * ub_pingpongNum + i));
        }
    }


}

